from sls_mqttClient import MQTTClient
import json

alltopics = ["/status/fly1/PR", "/status/fly2/PR","/status/fly1/runingTask"]
class UAVMirror:
    def __init__(self,idstr:str = None):
        #字符串形式的id
        self.idstr:str = idstr
        #位置
        self.GNSS_position ={"latitude":0,"longitude":0}
        self.NED_position = {"N":0,"E":0,"D":0}

        self.NED_velocity = {"N":0,"E":0,"D":0}
        #相对高度
        self.relative_altitude = {"NED":0,"msl":0,"ell":0,"laser":0}
        #偏航角
        self.vehicle_yaw = 0

        self.takeoff_NED_position = {"N":0,"E":0,"D":0}
        self.takeoff_GNSS_position = {"latitude":0,"longitude":0}
        self.takeoff_altitude = 0

        #状态
        self.is_flying_status = False
        self.is_arming_status = False

        #姿态
        self.q_FRD_REFTO_NED = [0,0,0,0] #w,x,y,z 

        self.runingTask = {
            "missionsType": None,
            "datasets": None
        }
        self.missionsIndex:int = None
        self.missionsList = []

        self.scheduleState = None
        self.nav_state =None


class sls_Agent:
    def __init__(self):
        self.uav_mirror_1 = UAVMirror("fly1")
        self.uav_mirror_2 = UAVMirror("fly2")
        self.uav_mirrors = [self.uav_mirror_1,self.uav_mirror_2]

        self.agent_update_callback = None
        self.agent_connect_callback = None
        self.runingTask_callback = None
        self.readList_callback = None

        self.mqttClient = MQTTClient()
        self.mqttClient.set_message_handle_callback(self.message_handle)
        self.mqttClient.set_connect_handle_callback(self.connect_handle)

    def find_mirror_by_idstr(self,idstr:str) ->UAVMirror:
        for mirror in self.uav_mirrors:
            if mirror.idstr == idstr:
                return mirror
        return None

    def set_agent_update_callback(self,callback):
        self.agent_update_callback = callback

    def set_agent_connect_callback(self,callback):
        self.agent_connect_callback = callback
    
    def set_runingTask_callback(self,callback):
        self.runingTask_callback = callback
    
    def set_readList_callback(self,callback):
        self.readList_callback = callback
        
    def mqtt_connect(self):
        self.mqttClient.connect()

    def message_handle(self, client, userdata, msg):
        topic = msg.topic
        message = msg.payload.decode()
        # 定义一个更新 uav_mirror 属性的函数
        def update_uav_mirror(uav_mirror, messageDict):
            if "GNSS_position" in messageDict:
                uav_mirror.GNSS_position = messageDict["GNSS_position"]
            if "NED_position" in messageDict:
                uav_mirror.NED_position = messageDict["NED_position"]
            if "NED_velocity" in messageDict:
                uav_mirror.NED_velocity = messageDict["NED_velocity"]
            if "relative_altitude" in messageDict:
                uav_mirror.relative_altitude = messageDict["relative_altitude"]
            if "vehicle_yaw" in messageDict:
                uav_mirror.vehicle_yaw = messageDict["vehicle_yaw"]
            if "takeoff_NED_position" in messageDict:
                uav_mirror.takeoff_NED_position = messageDict["takeoff_NED_position"]
            if "takeoff_GNSS_position" in messageDict:
                uav_mirror.takeoff_GNSS_position = messageDict["takeoff_GNSS_position"]
            if "takeoff_altitude" in messageDict:
                uav_mirror.takeoff_altitude = messageDict["takeoff_altitude"]
            if "flying_status" in messageDict:
                uav_mirror.is_flying_status = messageDict["flying_status"]
            if "arming_status" in messageDict:
                uav_mirror.is_arming_status = messageDict["arming_status"]
            if "q_FRD_REFTO_NED" in messageDict:
                uav_mirror.q_FRD_REFTO_NED = messageDict["q_FRD_REFTO_NED"]
            if "nav_state" in messageDict:
                uav_mirror.nav_state = messageDict["nav_state"]

        def update_uav_mirror_runingTask(uav_mirror, messageDict):
            if "missionsType" in messageDict:
                uav_mirror.runingTask["missionsType"] = messageDict["missionsType"]
            if "datasets" in messageDict:
                uav_mirror.runingTask["datasets"] = messageDict["datasets"]
            if "missionsIndex" in messageDict:
                uav_mirror.missionsIndex = messageDict["missionsIndex"]
            # 其他属性的更新逻辑...
        # 判断 topic 是否为 "/status/fly1/PR" 或 "/status/fly2/PR"
        if topic in ["/status/fly1/PR", "/status/fly2/PR"]:
            messageDict = json.loads(message)
            if topic == "/status/fly1/PR":
                update_uav_mirror(self.uav_mirror_1, messageDict)
            elif topic == "/status/fly2/PR":
                update_uav_mirror(self.uav_mirror_2, messageDict)
            if self.agent_update_callback:
                self.agent_update_callback(client, userdata, msg)
        
        if topic in ["/status/fly1/runingTask","/status/fly2/runingTask"]: 
            messageDict = json.loads(message)
            if topic == "/status/fly1/runingTask":
                update_uav_mirror_runingTask(self.uav_mirror_1, messageDict)
            elif topic == "/status/fly2/runingTask":
                update_uav_mirror_runingTask(self.uav_mirror_2, messageDict)
            if self.runingTask_callback:
                self.runingTask_callback(client, userdata, msg)
        
        if topic in ["/status/fly1/missionsList","/status/fly2/missionsList"]:
            messageList = json.loads(message)
            
            if topic == "/status/fly1/missionsList":
                self.uav_mirror_1.missionsList.clear()
                for missionDict in messageList:
                    self.uav_mirror_1.missionsList.append(missionDict)
            elif topic == "/status/fly2/missionsList":
                self.uav_mirror_2.missionsList.clear()
                for missionDict in messageList:
                    self.uav_mirror_2.missionsList.append(missionDict)
            if self.readList_callback:
                self.readList_callback(client, userdata, msg)

        if topic in ["/status/fly1/scheduleState","/status/fly2/scheduleState"]:
            messageDict = json.loads(message)
            if topic == "/status/fly1/scheduleState":
                self.uav_mirror_1.missionsIndex = messageDict["index"]
                self.uav_mirror_1.scheduleState = messageDict["state"]
            elif topic == "/status/fly2/scheduleState":
                self.uav_mirror_2.missionsIndex = messageDict["index"]
                self.uav_mirror_2.scheduleState = messageDict["state"]

                

    def connect_handle(self, client, userdata, flags, rc):
        self.mqttClient.subscribe_topics(["/status/fly1/PR"])
        self.mqttClient.subscribe_topics(["/status/fly2/PR"])
        self.mqttClient.subscribe_topics(["/status/fly1/runingTask"])
        self.mqttClient.subscribe_topics(["/status/fly2/runingTask"])
        self.mqttClient.subscribe_topics(["/status/fly1/missionsList"])
        self.mqttClient.subscribe_topics(["/status/fly2/missionsList"])
        self.mqttClient.subscribe_topics(["/status/fly1/scheduleState"])
        self.mqttClient.subscribe_topics(["/status/fly2/scheduleState"])
        if self.agent_connect_callback:
            self.agent_connect_callback()

    def publish_MS_command(self,fly,command:str):
        if fly in ["fly1","fly2"]:
            if command in ["start","stop","read_list","restart","read_schedule_state"]:
                commandDict = {"method": command}
                self.mqttClient.publish_message(f"/command/{fly}/MS", commandDict)
    
    def publish_MS_appendMission(self,fly,missionDict):
        msgDict = {"method": "append_mission","mission":missionDict}
        self.mqttClient.publish_message(f"/command/{fly}/MS", msgDict)
    
    def publish_MS_addMission(self,fly,missionDict,index:int):
        msgDict = {"method": "add_mission","mission":missionDict,"index":index}
        self.mqttClient.publish_message(f"/command/{fly}/MS", msgDict)

    def publish_MS_deleteMission(self,fly,index:int):
        msgDict = {"method": "delete_mission","index":index}
        self.mqttClient.publish_message(f"/command/{fly}/MS", msgDict)
            
    def fly1_start_Ms(self):
        self.publish_MS_command("fly1","start")

    def fly1_stop_Ms(self):
        self.publish_MS_command("fly1","stop")

    def fly2_start_Ms(self):
        self.publish_MS_command("fly2","start")

    def fly2_stop_Ms(self):
        self.publish_MS_command("fly2","stop")

    def fly1_restart_Ms(self):
        self.publish_MS_command("fly1","restart")
    
    def fly2_restart_Ms(self):
        self.publish_MS_command("fly2","restart")
    
    def ask_missionslist(self,fly:str):
        self.publish_MS_command(fly,"read_list")

    def publish_PX_command(self,fly,command:str):
        if fly in ["fly1","fly2"]:
            if command in ["set_Offboard"]:
                commandDict = {"method": command}
                self.mqttClient.publish_message(f"/command/{fly}/PX", commandDict)
    
    def fly1_set_Offboard(self):
        self.publish_PX_command("fly1","set_Offboard")

    





